
dot_motion_shader = [
    """ #version 140

        uniform sampler2D p3d_Texture0;
        uniform mat4 p3d_ModelViewProjectionMatrix;

        in vec4 p3d_Vertex;
        in vec2 p3d_MultiTexCoord0;

        out vec4 color;

        void main(void) {
            vec4 newvertex;
            float dot_i;
            float dot_scale = 0.01;
            float dot_x, dot_y, dot_color;
            float maxi = 10000.0;
            vec4 dot_properties;

            dot_i = float(p3d_Vertex[1]);
            dot_properties = texture2D(p3d_Texture0, vec2(dot_i/maxi, 0.0));

            dot_x = dot_properties[2];
            dot_y = dot_properties[1];
            dot_color = dot_properties[0];

            newvertex = p3d_Vertex;

            if (dot_x*dot_x + dot_y*dot_y > 1.0*1.0 || dot_i > 1200) { // only plot 1200 dots inside the circle
                newvertex[0] = 0.0;
                newvertex[1] = 0.0;
                newvertex[2] = 0.0;
            } else {
                newvertex[0] = p3d_Vertex[0]*dot_scale+dot_x;
                newvertex[1] = 0;
                newvertex[2] = p3d_Vertex[2]*dot_scale+dot_y;
            }

            color = vec4(dot_color, dot_color, dot_color, 1.0);

            gl_Position = p3d_ModelViewProjectionMatrix * newvertex;
        }
    """,

    """ #version 140
        in vec4 color;
        out vec4 gl_FragColor;

        void main() {
            gl_FragColor = color;
        }

    """
]

if __name__ == "__main__":
    import sys
    import numpy as np
    import os

    sys.path.append("../modules")

    from shared import Shared

    #coherences = [50,100]
    coherences = [0,25,50,100] # Comment line 66 and uncomment this line
    velocity = 0.13636
    orientations = [-90., 90.]
    stimuli = []

    for orientation in orientations:
        for coherence in coherences:
            stimuli.append([coherence, coherence,
                        orientation, orientation,
                        velocity, velocity])

    t_stimuli = [20] * len(stimuli)
    #stimuli = [[100,100,90,-90,0.13636,0.13636], [100,100,-90,90,0.13636,0.13636]]
    shared = Shared(r'D:\kumaresh_free_swimming_4fish_data', "dot_hanna_12_08_2022_5mgul_cortisol", stimuli = stimuli, t_stimuli = t_stimuli, shaders=[dot_motion_shader], python_file = __file__)

    #stimuli.append([0,0,0,0,0,0]) # baseline


    shared.start_threads()


    from scene import Scene
    from panda3d.core import *

    class World(Scene):

        def __init__(self):

            # Init the scene
            Scene.__init__(self, shared=shared)

            self.background_circle = dict({})
            self.circles = dict({})
            self.dummytex = dict({})
            self.coherent_change_vector_ind = dict({})

            self.dots_position = dict({})

            self.dots_circular_d = dict({})
            self.dots_circular_ang = dict({})
            self.dots_x = dict({})
            self.dots_y = dict({})

            for fish_index in range(4):
                self.background_circle[fish_index] = self.create_circles(1, edges=30)
                self.background_circle[fish_index].reparentTo(self.fish_nodes[fish_index])

                self.background_circle[fish_index].setScale(1, 0, 1)
                self.background_circle[fish_index].setColor(0,0, 0)
                self.background_circle[fish_index].setPos(0, 0.001, 0)

                filepath = os.path.join(os.path.split(__file__)[0], "circles.bam")
                self.circles[fish_index] = self.loader.loadModel(Filename.fromOsSpecific(filepath))
                self.circles[fish_index].reparentTo(self.fish_nodes[fish_index])

                self.dummytex[fish_index] = Texture("dummy texture")
                self.dummytex[fish_index].setup2dTexture(10000, 1, Texture.T_float, Texture.FRgb32)
                self.dummytex[fish_index].setMagfilter(Texture.FTNearest)

                ts1 = TextureStage("part2")
                ts1.setSort(-100)

                self.circles[fish_index].setTexture(ts1, self.dummytex[fish_index])
                self.circles[fish_index].setShader(self.compiled_shaders[0])

                self.dots_circular_d[fish_index] = np.random.random(10000).astype(np.float32)
                self.dots_circular_ang[fish_index] = 360 * np.random.random(10000).astype(np.float32)

                self.dots_x[fish_index] = 2 * np.random.random(10000).astype(np.float32) - 1
                self.dots_y[fish_index] = 2 * np.random.random(10000).astype(np.float32) - 1

                self.dots_position[fish_index] = np.empty((1, 10000, 3)).astype(np.float32)
                self.dots_position[fish_index][0, :, 0] = self.dots_x[fish_index]
                self.dots_position[fish_index][0, :, 1] = self.dots_y[fish_index]
                self.dots_position[fish_index][0, :, 2] = np.random.randint(0, 3, 10000).astype(np.float32) * 0 + 1  # 0 will be black, 1, will be white, 2 will be hidden

                memoryview(self.dummytex[fish_index].modify_ram_image())[:] = self.dots_position[fish_index].tobytes()

        def stimulus(self, fish_index, i_stimulus, t_stimulus, frame, time, dt, trial):

            #i_stimulus = 0
            coherence_left, coherence_right, orientation_left, orientation_right, velocity_left, velocity_right = stimuli[i_stimulus]

            fish_angle = self.shared.current_fish_accumulated_orientation_lowpass[fish_index].value
            #fish_angle = 0

            x = self.shared.current_fish_position_x[fish_index].value
            y = self.shared.current_fish_position_y[fish_index].value
            #x = 0.; y = 0.

            slope = np.tan((fish_angle) * np.pi / 180.)
            diff = 0.0 #diff = 0.05 * np.sqrt(1 + slope ** 2)
            target = slope * (self.dots_x[fish_index] - x) + y

            if np.cos(fish_angle * np.pi / 180.) > 0:
                target_right = target - diff
                target_left = target + diff

                ind_right = self.dots_y[fish_index] < target_right
                ind_left = self.dots_y[fish_index] > target_left
                flag = 0
            else:
                target_right = target + diff
                target_left = target - diff

                ind_right = self.dots_y[fish_index] > target_right
                ind_left = self.dots_y[fish_index] < target_left
                flag = 1

            if frame == 0:
                random_vector = np.random.randint(100, size=10000)

                self.coherent_change_vector_ind[fish_index] = random_vector < coherence_left

            if 5 <= time < 15:

                self.dots_x[fish_index][self.coherent_change_vector_ind[fish_index] & ind_right] += (np.cos((-orientation_right + fish_angle) * np.pi / 180) ) * \
                                                                                            velocity_right * dt
                self.dots_y[fish_index][self.coherent_change_vector_ind[fish_index] & ind_right] += (np.sin((-orientation_right + fish_angle) * np.pi / 180)) * \
                                                                                            velocity_right * dt

                self.dots_x[fish_index][self.coherent_change_vector_ind[fish_index] & ind_left] += (np.cos((-orientation_left + fish_angle) * np.pi / 180)) * \
                                                                                           velocity_left * dt
                self.dots_y[fish_index][self.coherent_change_vector_ind[fish_index] & ind_left] += (np.sin((-orientation_left + fish_angle) * np.pi / 180)) * \
                                                                                           velocity_left * dt

            # everything gets randomly redraw
            k = np.random.random(10000)
            ind = np.where(k < dt / 0.2)[0]

            self.dots_x[fish_index][ind] = 2 * np.random.random(len(ind)).astype(np.float32) - 1  # x
            self.dots_y[fish_index][ind] = 2 * np.random.random(len(ind)).astype(np.float32) - 1  # y

            # Wrap them
            self.dots_x[fish_index] = (self.dots_x[fish_index] + 1) % 2 - 1
            self.dots_y[fish_index] = (self.dots_y[fish_index] + 1) % 2 - 1

            self.dots_position[fish_index][0, :, 0] = self.dots_x[fish_index]
            self.dots_position[fish_index][0, :, 1] = self.dots_y[fish_index]

            if flag == 0:
                ind_hidden = (self.dots_y[fish_index] < target_right) | (self.dots_y[fish_index] > target_left)
            else:
                ind_hidden = (self.dots_y[fish_index] > target_right) | (self.dots_y[fish_index] < target_left)

            self.dots_position[fish_index][0, :, 2] = ind_hidden * 2


            memoryview(self.dummytex[fish_index].modify_ram_image())[:] = self.dots_position[fish_index].tobytes()

    try:
        world = World()
        world.run()
    except:
        shared.error("scene_error.txt", sys.exc_info())
